{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "9de31cb1",
      "metadata": {},
      "source": [
        "# Hidden Markov Model Example\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "colab_button",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/python/gtsam/examples/HMMExample.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "license_cell",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\nAtlanta, Georgia 30332-0415\nAll Rights Reserved\n\nAuthors: Frank Dellaert, et al. (see THANKS for the full author list)\n\nSee LICENSE for the license information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "colab_import",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n    import google.colab\n    %pip install --quiet gtsam-develop\nexcept ImportError:\n    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a18a7e9a",
      "metadata": {},
      "source": [
        "This example demonstrates a Hidden Markov Model (HMM) using discrete Bayes nets.\n",
        "\n",
        "A Hidden Markov Model consists of:\n",
        "- A sequence of hidden states over time\n",
        "- Transition probabilities between states\n",
        "- Observation/measurement probabilities for each state\n",
        "\n",
        "In this example:\n",
        "- We have 4 nodes (time steps) with 3 possible states each\n",
        "- We define a transition model that favors staying in the same state\n",
        "- We add measurements/observations at some (not all) time steps"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "7a95c00f",
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "import gtsam\n",
        "from gtsam import (DiscreteBayesNet, DiscreteFactorGraph, DiscreteMarginals,\n",
        "                   DiscreteKeys, Ordering)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "afcb6106",
      "metadata": {},
      "source": [
        "## Model Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "4a2bb3a7",
      "metadata": {},
      "outputs": [],
      "source": [
        "nr_nodes = 4\n",
        "nr_states = 3\n",
        "\n",
        "# Define variables as well as ordering\n",
        "ordering = Ordering()\n",
        "keys = []\n",
        "for k in range(nr_nodes):\n",
        "    key_i = (k, nr_states)\n",
        "    keys.append(key_i)\n",
        "    ordering.push_back(k)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5291abaa",
      "metadata": {},
      "source": [
        "## Create HMM as a DiscreteBayesNet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "598f34f2",
      "metadata": {},
      "outputs": [],
      "source": [
        "hmm = DiscreteBayesNet()\n",
        "\n",
        "# Define backbone transition model\n",
        "# \"8/1/1 1/8/1 1/1/8\" means the transition matrix favors staying in the same state\n",
        "# For each parent state (0,1,2), it gives probabilities for child states (0,1,2)\n",
        "transition = \"8/1/1 1/8/1 1/1/8\"\n",
        "\n",
        "# Add transition conditionals for k=1,2,3 conditioned on k-1\n",
        "for k in range(1, nr_nodes):\n",
        "    # Helper function to create DiscreteKeys for parents\n",
        "    parents = DiscreteKeys()\n",
        "    parents.push_back(keys[k - 1])\n",
        "    hmm.add(keys[k], parents, transition)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "a02a6369",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HMM\n",
            " \n",
            "size: 6\n",
            "conditional 0:  P( 1 | 0 ):\n",
            " Choice(1) \n",
            " 0 Choice(0) \n",
            " 0 0 Leaf  0.8\n",
            " 0 1 Leaf  0.1\n",
            " 0 2 Leaf  0.1\n",
            " 1 Choice(0) \n",
            " 1 0 Leaf  0.1\n",
            " 1 1 Leaf  0.8\n",
            " 1 2 Leaf  0.1\n",
            " 2 Choice(0) \n",
            " 2 0 Leaf  0.1\n",
            " 2 1 Leaf  0.1\n",
            " 2 2 Leaf  0.8\n",
            "\n",
            "conditional 1:  P( 2 | 1 ):\n",
            " Choice(2) \n",
            " 0 Choice(1) \n",
            " 0 0 Leaf  0.8\n",
            " 0 1 Leaf  0.1\n",
            " 0 2 Leaf  0.1\n",
            " 1 Choice(1) \n",
            " 1 0 Leaf  0.1\n",
            " 1 1 Leaf  0.8\n",
            " 1 2 Leaf  0.1\n",
            " 2 Choice(1) \n",
            " 2 0 Leaf  0.1\n",
            " 2 1 Leaf  0.1\n",
            " 2 2 Leaf  0.8\n",
            "\n",
            "conditional 2:  P( 3 | 2 ):\n",
            " Choice(3) \n",
            " 0 Choice(2) \n",
            " 0 0 Leaf  0.8\n",
            " 0 1 Leaf  0.1\n",
            " 0 2 Leaf  0.1\n",
            " 1 Choice(2) \n",
            " 1 0 Leaf  0.1\n",
            " 1 1 Leaf  0.8\n",
            " 1 2 Leaf  0.1\n",
            " 2 Choice(2) \n",
            " 2 0 Leaf  0.1\n",
            " 2 1 Leaf  0.1\n",
            " 2 2 Leaf  0.8\n",
            "\n",
            "conditional 3:  P( 0 ):\n",
            " Choice(0) \n",
            " 0 Leaf  0.7\n",
            " 1 Leaf  0.2\n",
            " 2 Leaf  0.1\n",
            "\n",
            "conditional 4:  P( 1 ):\n",
            " Choice(1) \n",
            " 0 Leaf  0.1\n",
            " 1 Leaf  0.9\n",
            " 2 Leaf    0\n",
            "\n",
            "conditional 5:  P( 3 ):\n",
            " Choice(3) \n",
            " 0 Leaf  0.5\n",
            " 1 Leaf  0.4\n",
            " 2 Leaf  0.1\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Add some measurements, not needed for all time steps!\n",
        "# These are observation likelihoods\n",
        "hmm.add(keys[0], \"7/2/1\")  # Measurement at time 0\n",
        "hmm.add(keys[1], \"1/9/0\")  # Measurement at time 1\n",
        "hmm.add(keys[3], \"5/4/1\")  # Measurement at time 3 (last node)\n",
        "\n",
        "# Print the HMM\n",
        "hmm.print(\"HMM\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4283a552",
      "metadata": {},
      "source": [
        "## Inference: Maximum Probable Explanation (MPE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "876931f4",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "mpe: (0, 1) (1, 1) (2, 1) (3, 1)\n"
          ]
        }
      ],
      "source": [
        "# Convert to factor graph\n",
        "factor_graph = DiscreteFactorGraph(hmm)\n",
        "\n",
        "# Do max-product to find the most probable explanation (MPE)\n",
        "mpe = factor_graph.optimize()\n",
        "print(\"mpe:\", end=\"\")\n",
        "for k in range(nr_nodes):\n",
        "    print(f\" ({k}, {mpe[k]})\", end=\"\")\n",
        "print()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f968dca5",
      "metadata": {},
      "source": [
        "## Elimination, Sampling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "d4fa3530",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Eliminated\n",
            " \n",
            "size: 4\n",
            "conditional 0:  P( 0 | 1 ):\n",
            " Choice(1) \n",
            " 0 Choice(0) \n",
            " 0 0 Leaf 0.94915254\n",
            " 0 1 Leaf 0.033898305\n",
            " 0 2 Leaf 0.016949153\n",
            " 1 Choice(0) \n",
            " 1 0 Leaf 0.29166667\n",
            " 1 1 Leaf 0.66666667\n",
            " 1 2 Leaf 0.041666667\n",
            " 2 Choice(0) \n",
            " 2 0 Leaf 0.41176471\n",
            " 2 1 Leaf 0.11764706\n",
            " 2 2 Leaf 0.47058824\n",
            "\n",
            "conditional 1:  P( 1 | 2 ):\n",
            " Choice(2) \n",
            " 0 Choice(1) \n",
            " 0 0 Leaf 0.68604651\n",
            " 0 1 Leaf 0.31395349\n",
            " 0 2 Leaf    0\n",
            " 1 Choice(1) \n",
            " 1 0 Leaf 0.033016228\n",
            " 1 1 Leaf 0.96698377\n",
            " 1 2 Leaf    0\n",
            " 2 Choice(1) \n",
            " 2 0 Leaf 0.21454545\n",
            " 2 1 Leaf 0.78545455\n",
            " 2 2 Leaf    0\n",
            "\n",
            "conditional 2:  P( 2 | 3 ):\n",
            " Choice(3) \n",
            " 0 Choice(2) \n",
            " 0 0 Leaf 0.72746497\n",
            " 0 1 Leaf 0.23618821\n",
            " 0 2 Leaf 0.036346815\n",
            " 1 Choice(2) \n",
            " 1 0 Leaf 0.045088145\n",
            " 1 1 Leaf 0.9368897\n",
            " 1 2 Leaf 0.018022151\n",
            " 2 Choice(2) \n",
            " 2 0 Leaf 0.14716578\n",
            " 2 1 Leaf 0.38224599\n",
            " 2 2 Leaf 0.47058824\n",
            "\n",
            "conditional 3:  P( 3 ):\n",
            " Choice(3) \n",
            " 0 Leaf 0.36536251\n",
            " 1 Leaf 0.58948629\n",
            " 2 Leaf 0.045151196\n",
            "\n",
            "\n",
            "10 samples:\n",
            "sample: (0, 0) (1, 1) (2, 1) (3, 1)\n",
            "sample: (0, 1) (1, 1) (2, 1) (3, 1)\n",
            "sample: (0, 0) (1, 0) (2, 0) (3, 0)\n",
            "sample: (0, 1) (1, 1) (2, 1) (3, 1)\n",
            "sample: (0, 0) (1, 1) (2, 1) (3, 1)\n",
            "sample: (0, 0) (1, 1) (2, 1) (3, 1)\n",
            "sample: (0, 0) (1, 0) (2, 0) (3, 0)\n",
            "sample: (0, 1) (1, 1) (2, 1) (3, 1)\n",
            "sample: (0, 1) (1, 1) (2, 1) (3, 1)\n",
            "sample: (0, 0) (1, 1) (2, 1) (3, 1)\n"
          ]
        }
      ],
      "source": [
        "# Create solver and eliminate\n",
        "# This will create a DAG ordered with arrow of time reversed\n",
        "chordal = factor_graph.eliminateSequential(ordering)\n",
        "chordal.print(\"Eliminated\\n\")\n",
        "\n",
        "# We can also sample from it\n",
        "print(\"\\n10 samples:\")\n",
        "for k in range(10):\n",
        "    sample = chordal.sample()\n",
        "    print(\"sample:\", end=\"\")\n",
        "    for i in range(nr_nodes):\n",
        "        print(f\" ({i}, {sample[i]})\", end=\"\")\n",
        "    print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "84f4200f",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Computing Node Marginals ..\n",
            "marginal 0[0.44714654 0.5170319  0.03582156]\n",
            "marginal 1[0.23647637 0.76352363 0.        ]\n",
            "marginal 2[0.29901199 0.65583682 0.0451512 ]\n",
            "marginal 3[0.36536251 0.58948629 0.0451512 ]\n"
          ]
        }
      ],
      "source": [
        "# Or compute the marginals. This re-eliminates the FG into a Bayes tree\n",
        "print(\"\\nComputing Node Marginals ..\")\n",
        "marginals = DiscreteMarginals(factor_graph)\n",
        "for k in range(nr_nodes):\n",
        "    marg_probs = marginals.marginalProbabilities((k, nr_states))\n",
        "    print(f\"marginal {k}{marg_probs}\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}